// Copyright(c) Sonosemi Medical Ltd. Embedded Software Team.
// Author: Zhang Yisong(zhangyisong@sonosemi.com)

#ifndef __ZPP_CALLBACK_HPP__
#define __ZPP_CALLBACK_HPP__

#include <cstddef>
#include <new>
#include <type_traits>
#include <utility>
#include <zpp/assert.hpp>
namespace zpp {

template <typename t_signature, size_t k_storage_size = 3 * sizeof(void *)>
class callback;

template <typename t_ret, typename... t_args, size_t k_storage_size>
class callback<t_ret(t_args...), k_storage_size> {
public:
  callback() = default;
  template <typename t_callable>
    requires std::is_invocable_r_v<t_ret, t_callable, t_args...>
  callback(t_callable &&callable) {
    static_assert(sizeof(t_callable) <= k_storage_size,
                  "Storage size too small for this lambda");
    auto *wrapper = reinterpret_cast<LambdaWrapper<t_callable> *>(&m_storage);
    new (wrapper) LambdaWrapper<t_callable>{std::forward<t_callable>(callable)};

    m_invoke = &LambdaWrapper<t_callable>::invoke;
    m_context = wrapper;
  }

  callback(t_ret (*func)(t_args...)) {
    static_assert(sizeof(t_ret(*)(t_args...)) <= k_storage_size,
                  "Storage size too small for function pointer");

    auto *wrapper = reinterpret_cast<FunctionWrapper *>(&m_storage);
    wrapper->func = func;

    m_invoke = &FunctionWrapper::invoke;
    m_context = wrapper;
  }
  auto reset() -> void {
    m_invoke = nullptr;
    m_context = nullptr;
  }
  auto operator()(t_args... args) const -> t_ret {
    return m_invoke(m_context, std::forward<t_args>(args)...);
  }
  template <typename T> callback(t_ret (T::*method)(t_args...), T &obj) {
    using MethodWrapper = MethodWrapper<T>;
    static_assert(sizeof(MethodWrapper) <= k_storage_size,
                  "Storage size too small for this method");

    auto *wrapper = reinterpret_cast<MethodWrapper *>(&m_storage);
    wrapper->method = method;
    wrapper->obj = &obj;

    m_invoke = &MethodWrapper::invoke;
    m_context = wrapper;
  }
  ~callback() { reset(); }

  auto empty() -> bool const { return m_invoke == nullptr; }

  callback(const callback &) = delete;
  callback &operator=(const callback &) = delete;
  callback(callback &&other) {
    m_invoke = other.m_invoke;
    m_context = other.m_context;

    other.m_invoke = nullptr;
    other.m_context = nullptr;
  }
  callback &operator=(callback &&other) {
    if (this != &other) {
      m_invoke = other.m_invoke;
      m_context = other.m_context;

      other.m_invoke = nullptr;
      other.m_context = nullptr;
    }
    return *this;
  }

private:
  using invoke_t = t_ret (*)(void *, t_args...);
  static auto dumb(void *, t_args...) -> t_ret {
    if constexpr (std::is_same_v<void, t_ret>) {
    } else {
      return t_ret{};
    }
  }
  template <typename Lambda> struct LambdaWrapper {
    Lambda lambda;
    explicit LambdaWrapper(Lambda &&l) : lambda(std::forward<Lambda>(l)) {}
    static auto invoke(void *ctx, t_args... args) -> t_ret {
      auto *self = reinterpret_cast<LambdaWrapper *>(ctx);
      return self->lambda(std::forward<t_args>(args)...);
    }
  };

  struct FunctionWrapper {
    t_ret (*func)(t_args...);

    static auto invoke(void *ctx, t_args... args) -> t_ret {
      auto *self = reinterpret_cast<FunctionWrapper *>(ctx);
      return self->func(std::forward<t_args>(args)...);
    }
  };

  template <typename T> struct MethodWrapper {
    t_ret (T::*method)(t_args...);
    T *obj;

    static auto invoke(void *ctx, t_args... args) -> t_ret {
      auto *self = reinterpret_cast<MethodWrapper *>(ctx);
      return (self->obj->*(self->method))(std::forward<t_args>(args)...);
    }
  };
  typename std::aligned_storage<k_storage_size>::type m_storage;
  invoke_t m_invoke{dumb};
  void *m_context = nullptr;
};

} // namespace zpp

#endif //__ZPP_CALLBACK_HPP__
